package org.cybzacg.encryption.impl;

import org.cybzacg.encryption.core.AbstractEncryptor;
import org.cybzacg.encryption.context.EncryptionContext;
import org.cybzacg.encryption.model.EncryptionResult;
import org.cybzacg.encryption.model.ValidationResult;
import org.cybzacg.encryption.enums.EncryptionAlgorithm;
import org.cybzacg.encryption.enums.EncryptionType;
import org.cybzacg.encryption.enums.NationalAlgorithm;
import org.cybzacg.encryption.utils.ValidationUtils;
import org.cybzacg.encryption.utils.PerformanceUtils;

import javax.crypto.*;
import java.security.*;
import java.security.interfaces.*;
import java.security.spec.*;
import java.util.Arrays;
import java.math.BigInteger;

/**
 * SM2加密器实现
 * 中国商用密码椭圆曲线公钥密码算法
 *
 * @author encryption team
 * @version 1.0
 */
public class SM2Encryptor extends AbstractEncryptor {

    private static final String ALGORITHM_NAME = "SM2";
    private static final int KEY_SIZE = 256; // SM2密钥长度256位
    private static final int SIGNATURE_SIZE = 64; // SM2签名长度64字节

    /**
     * 构造函数
     */
    public SM2Encryptor() {
        super(EncryptionAlgorithm.SM2);
    }

    @Override
    protected byte[] doEncrypt(byte[] data, EncryptionContext context) {
        try (PerformanceUtils.PerformanceContext performanceContext =
                PerformanceUtils.startMonitoring("SM2 encryption")) {

            // 验证输入参数
            ValidationResult validationResult = validateEncryptionInput(data, context);
            if (!validationResult.isValid()) {
                throw new IllegalArgumentException(validationResult.getMessage());
            }

            // 获取公钥
            PublicKey publicKey = getPublicKey(context);
            if (publicKey == null) {
                throw new IllegalArgumentException("SM2 public key not available in context");
            }

            // 验证公钥
            validationResult = validatePublicKey(publicKey);
            if (!validationResult.isValid()) {
                throw new IllegalArgumentException(validationResult.getMessage());
            }

            // 执行SM2加密
            return performSM2Encryption(data, publicKey);

        } catch (Exception e) {
            throw new RuntimeException("SM2 encryption failed: " + e.getMessage(), e);
        }
    }

    @Override
    protected byte[] doDecrypt(byte[] encryptedData, EncryptionContext context) {
        try (PerformanceUtils.PerformanceContext performanceContext =
                PerformanceUtils.startMonitoring("SM2 decryption")) {

            // 验证输入参数
            ValidationResult validationResult = validateDecryptionInput(encryptedData, context);
            if (!validationResult.isValid()) {
                throw new IllegalArgumentException(validationResult.getMessage());
            }

            // 获取私钥
            PrivateKey privateKey = getPrivateKey(context);
            if (privateKey == null) {
                throw new IllegalArgumentException("SM2 private key not available in context");
            }

            // 验证私钥
            validationResult = validatePrivateKey(privateKey);
            if (!validationResult.isValid()) {
                throw new IllegalArgumentException(validationResult.getMessage());
            }

            // 执行SM2解密
            return performSM2Decryption(encryptedData, privateKey);

        } catch (Exception e) {
            throw new RuntimeException("SM2 decryption failed: " + e.getMessage(), e);
        }
    }

    @Override
    protected void doInitialize(EncryptionContext context) {
        // SM2特定的初始化逻辑
        if (context != null) {
            // 添加SM2特定的初始化参数
            context.setAttribute("algorithm", ALGORITHM_NAME);
            context.setAttribute("keySize", KEY_SIZE);
            context.setAttribute("signatureSize", SIGNATURE_SIZE);
        }
    }

    @Override
    protected ValidationResult doValidate(EncryptionContext context) {
        if (context == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Context cannot be null")
                    .build();
        }

        // 验证公钥
        PublicKey publicKey = getPublicKey(context);
        if (publicKey != null) {
            ValidationResult result = validatePublicKey(publicKey);
            if (!result.isValid()) {
                return result;
            }
        }

        // 验证私钥
        PrivateKey privateKey = getPrivateKey(context);
        if (privateKey != null) {
            ValidationResult result = validatePrivateKey(privateKey);
            if (!result.isValid()) {
                return result;
            }
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM2 context validation passed")
                .build();
    }

    @Override
    protected boolean doSupports(EncryptionContext context) {
        if (context == null) {
            return false;
        }

        // 检查算法匹配
        if (getAlgorithm() != context.getAlgorithm()) {
            return false;
        }

        // 检查加密类型匹配
        if (getType() != context.getEncryptionType()) {
            return false;
        }

        return true;
    }


    /**
     * SM2数字签名
     *
     * @param data 待签名数据
     * @param privateKey 私钥
     * @return 签名值（64字节）
     * @throws Exception 签名异常
     */
    public static byte[] sign(byte[] data, PrivateKey privateKey) throws Exception {
        if (data == null) {
            throw new IllegalArgumentException("Data to sign cannot be null");
        }

        if (privateKey == null) {
            throw new IllegalArgumentException("Private key cannot be null");
        }

        try {
            Signature signature = Signature.getInstance("SM3withSM2");
            signature.initSign(privateKey);
            signature.update(data);
            return signature.sign();
        } catch (Exception e) {
            throw new Exception("SM2 signing failed: " + e.getMessage(), e);
        }
    }

    /**
     * SM2数字签名验证
     *
     * @param data 原始数据
     * @param signature 签名值
     * @param publicKey 公钥
     * @return 验证结果
     */
    public static boolean verifySignature(byte[] data, byte[] signature, PublicKey publicKey) {
        if (data == null) {
            throw new IllegalArgumentException("Data to verify cannot be null");
        }

        if (signature == null) {
            throw new IllegalArgumentException("Signature cannot be null");
        }

        if (publicKey == null) {
            throw new IllegalArgumentException("Public key cannot be null");
        }

        try {
            Signature sig = Signature.getInstance("SM3withSM2");
            sig.initVerify(publicKey);
            sig.update(data);
            return sig.verify(signature);
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 生成SM2密钥对
     *
     * @return SM2密钥对
     * @throws Exception 生成异常
     */
    public static KeyPair generateKeyPair() throws Exception {
        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("SM2");
            keyGen.initialize(KEY_SIZE);
            return keyGen.generateKeyPair();
        } catch (Exception e) {
            throw new Exception("Failed to generate SM2 key pair: " + e.getMessage(), e);
        }
    }

    /**
     * 从字节数组恢复SM2公钥
     *
     * @param encoded 公钥编码字节数组
     * @return SM2公钥
     * @throws Exception 恢复异常
     */
    public static PublicKey restorePublicKey(byte[] encoded) throws Exception {
        try {
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encoded);
            KeyFactory keyFactory = KeyFactory.getInstance("SM2");
            return keyFactory.generatePublic(keySpec);
        } catch (Exception e) {
            throw new Exception("Failed to restore SM2 public key: " + e.getMessage(), e);
        }
    }

    /**
     * 从字节数组恢复SM2私钥
     *
     * @param encoded 私钥编码字节数组
     * @return SM2私钥
     * @throws Exception 恢复异常
     */
    public static PrivateKey restorePrivateKey(byte[] encoded) throws Exception {
        try {
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(encoded);
            KeyFactory keyFactory = KeyFactory.getInstance("SM2");
            return keyFactory.generatePrivate(keySpec);
        } catch (Exception e) {
            throw new Exception("Failed to restore SM2 private key: " + e.getMessage(), e);
        }
    }

    /**
     * 验证加密输入参数
     */
    private ValidationResult validateEncryptionInput(byte[] data, EncryptionContext context) {
        // 验证数据
        if (data == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM2 encryption data cannot be null")
                    .build();
        }

        if (data.length == 0) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM2 encryption data cannot be empty")
                    .build();
        }

        // 验证数据长度
        if (data.length > getMaxDataSize()) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("SM2 encryption data size %d bytes exceeds maximum %d bytes",
                        data.length, getMaxDataSize()))
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM2 encryption input is valid")
                .build();
    }

    /**
     * 验证解密输入参数
     */
    private ValidationResult validateDecryptionInput(byte[] encryptedData, EncryptionContext context) {
        // 验证加密数据
        if (encryptedData == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM2 encrypted data cannot be null")
                    .build();
        }

        if (encryptedData.length == 0) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM2 encrypted data cannot be empty")
                    .build();
        }

        // 验证加密数据长度
        if (encryptedData.length > getMaxDataSize()) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("SM2 encrypted data size %d bytes exceeds maximum %d bytes",
                        encryptedData.length, getMaxDataSize()))
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM2 decryption input is valid")
                .build();
    }

    /**
     * 验证公钥
     */
    public ValidationResult validatePublicKey(PublicKey publicKey) {
        if (publicKey == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM2 public key cannot be null")
                    .build();
        }

        // 检查密钥算法
        if (!"SM2".equals(publicKey.getAlgorithm())) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Public key algorithm must be SM2")
                    .build();
        }

        // 检查密钥格式
        if (!(publicKey instanceof ECPublicKey)) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Public key must be ECPublicKey")
                    .build();
        }

        ECPublicKey ecPublicKey = (ECPublicKey) publicKey;

        // 检查曲线参数
        ECParameterSpec params = ecPublicKey.getParams();
        if (params == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("EC parameters cannot be null")
                    .build();
        }

        // 检查曲线阶数
        BigInteger order = params.getOrder();
        if (order == null || order.bitLength() != KEY_SIZE) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message(String.format("EC curve order must be %d bits", KEY_SIZE))
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM2 public key is valid")
                .build();
    }

    /**
     * 验证私钥
     */
    public ValidationResult validatePrivateKey(PrivateKey privateKey) {
        if (privateKey == null) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM2 private key cannot be null")
                    .build();
        }

        // 检查密钥算法
        if (!"SM2".equals(privateKey.getAlgorithm())) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Private key algorithm must be SM2")
                    .build();
        }

        // 检查密钥格式
        if (!(privateKey instanceof ECPrivateKey)) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Private key must be ECPrivateKey")
                    .build();
        }

        ECPrivateKey ecPrivateKey = (ECPrivateKey) privateKey;

        // 检查私钥值
        BigInteger s = ecPrivateKey.getS();
        if (s == null || s.compareTo(BigInteger.ONE) < 0 ||
            s.compareTo(ecPrivateKey.getParams().getOrder().subtract(BigInteger.ONE)) >= 0) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("Private key scalar is out of valid range")
                    .build();
        }

        return new ValidationResult.Builder()
                .valid(true)
                .message("SM2 private key is valid")
                .build();
    }

    /**
     * 获取公钥
     */
    private PublicKey getPublicKey(EncryptionContext context) {
        // 尝试从上下文中获取公钥
        Object publicKeyObj = context.getAttribute("publicKey");
        if (publicKeyObj instanceof PublicKey) {
            return (PublicKey) publicKeyObj;
        }

        // 尝试从上下文中获取公钥字节数组
        Object publicKeyBytesObj = context.getAttribute("publicKeyBytes");
        if (publicKeyBytesObj instanceof byte[]) {
            try {
                return restorePublicKey((byte[]) publicKeyBytesObj);
            } catch (Exception e) {
                return null;
            }
        }

        return null;
    }

    /**
     * 获取私钥
     */
    private PrivateKey getPrivateKey(EncryptionContext context) {
        // 尝试从上下文中获取私钥
        Object privateKeyObj = context.getAttribute("privateKey");
        if (privateKeyObj instanceof PrivateKey) {
            return (PrivateKey) privateKeyObj;
        }

        // 尝试从上下文中获取私钥字节数组
        Object privateKeyBytesObj = context.getAttribute("privateKeyBytes");
        if (privateKeyBytesObj instanceof byte[]) {
            try {
                return restorePrivateKey((byte[]) privateKeyBytesObj);
            } catch (Exception e) {
                return null;
            }
        }

        return null;
    }

    /**
     * 执行SM2加密
     */
    private byte[] performSM2Encryption(byte[] data, PublicKey publicKey) throws Exception {
        // 简化的SM2加密实现
        // 实际应用中需要完整的SM2加密协议
        // 这里使用ECIES作为示例实现

        try {
            // 创建SM2加密器
            Cipher cipher = Cipher.getInstance("SM2");
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);

            // 执行加密
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new Exception("SM2 encryption operation failed: " + e.getMessage(), e);
        }
    }

    /**
     * 执行SM2解密
     */
    private byte[] performSM2Decryption(byte[] encryptedData, PrivateKey privateKey) throws Exception {
        // 简化的SM2解密实现
        // 实际应用中需要完整的SM2解密协议

        try {
            // 创建SM2解密器
            Cipher cipher = Cipher.getInstance("SM2");
            cipher.init(Cipher.DECRYPT_MODE, privateKey);

            // 执行解密
            return cipher.doFinal(encryptedData);
        } catch (Exception e) {
            throw new Exception("SM2 decryption operation failed: " + e.getMessage(), e);
        }
    }

    /**
     * 获取最大数据大小
     */
    private int getMaxDataSize() {
        // SM2最大加密数据大小限制
        return 1024; // 1KB
    }

    /**
     * 检查SM2算法是否可用
     *
     * @return 是否可用
     */
    public static boolean isAvailable() {
        try {
            KeyPairGenerator.getInstance("SM2");
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 获取算法信息
     *
     * @return 算法信息
     */
    public static String getAlgorithmInfo() {
        return String.format(
            "SM2 Algorithm Information:\n" +
            "  Algorithm: %s\n" +
            "  Key Size: %d bits\n" +
            "  Signature Size: %d bytes\n" +
            "  Available: %s",
            ALGORITHM_NAME, KEY_SIZE, SIGNATURE_SIZE, isAvailable()
        );
    }

    /**
     * SM2密钥对生成器
     */
    public static class SM2KeyPairGenerator {

        /**
         * 生成随机SM2密钥对
         *
         * @return SM2密钥对
         * @throws Exception 生成异常
         */
        public static KeyPair generateRandomKeyPair() throws Exception {
            return SM2Encryptor.generateKeyPair();
        }

        /**
         * 从种子生成SM2密钥对
         *
         * @param seed 种子
         * @return SM2密钥对
         * @throws Exception 生成异常
         */
        public static KeyPair generateKeyPairFromSeed(byte[] seed) throws Exception {
            try {
                SecureRandom random = SecureRandom.getInstance("SHA1PRNG");
                random.setSeed(seed);

                KeyPairGenerator keyGen = KeyPairGenerator.getInstance("SM2");
                keyGen.initialize(KEY_SIZE, random);
                return keyGen.generateKeyPair();
            } catch (Exception e) {
                throw new Exception("Failed to generate SM2 key pair from seed: " + e.getMessage(), e);
            }
        }

        /**
         * 生成指定参数的SM2密钥对
         *
         * @param params 椭圆曲线参数
         * @return SM2密钥对
         * @throws Exception 生成异常
         */
        public static KeyPair generateKeyPair(ECParameterSpec params) throws Exception {
            try {
                KeyPairGenerator keyGen = KeyPairGenerator.getInstance("SM2");
                keyGen.initialize(params);
                return keyGen.generateKeyPair();
            } catch (Exception e) {
                throw new Exception("Failed to generate SM2 key pair with params: " + e.getMessage(), e);
            }
        }
    }

    /**
     * SM2签名工具
     */
    public static class SM2Signer {

        private final PrivateKey privateKey;

        /**
         * 构造函数
         *
         * @param privateKey 私钥
         */
        public SM2Signer(PrivateKey privateKey) {
            this.privateKey = privateKey;
        }

        /**
         * 签名数据
         *
         * @param data 待签名数据
         * @return 签名值
         * @throws Exception 签名异常
         */
        public byte[] sign(byte[] data) throws Exception {
            return SM2Encryptor.sign(data, privateKey);
        }

        /**
         * 签名字符串
         *
         * @param data 待签名字符串
         * @return 签名值
         * @throws Exception 签名异常
         */
        public byte[] sign(String data) throws Exception {
            return sign(data.getBytes(java.nio.charset.StandardCharsets.UTF_8));
        }

        /**
         * 验证签名
         *
         * @param data 原始数据
         * @param signature 签名值
         * @param publicKey 公钥
         * @return 验证结果
         */
        public boolean verify(byte[] data, byte[] signature, PublicKey publicKey) {
            return SM2Encryptor.verifySignature(data, signature, publicKey);
        }

        /**
         * 验证签名
         *
         * @param data 原始字符串
         * @param signature 签名值
         * @param publicKey 公钥
         * @return 验证结果
         */
        public boolean verify(String data, byte[] signature, PublicKey publicKey) {
            return verify(data.getBytes(java.nio.charset.StandardCharsets.UTF_8), signature, publicKey);
        }
    }
}
